package org.example.loadbalancer;

import org.example.model.ServiceMetaInfo;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;

/**
 * @description 一致性哈希算法负载均衡算法
 * @author sats@jz
 * @date 2025年01月13日 14:32
 */
public class ConsistentHashLoadBalancer implements LoadBalancer{

    /*
     * 一致性Hash还
    */
    private final TreeMap<Integer, ServiceMetaInfo> circle = new TreeMap<>();

    /*
     * 虚拟节点数量
    */
    private static final int VIRTUAL_NODES = 100;

    /*
     * 添加虚拟节点
    */
    public void addVirtualNode(ServiceMetaInfo serviceMetaInfo){
        for (int i = 0; i < VIRTUAL_NODES; i++) {
            int hash = getHash(serviceMetaInfo.getServiceAddress() + "#" + i);
            circle.put(hash, serviceMetaInfo);
        }
    }

    /*
     * 移除虚拟节点
    */

    public void removeVirtualNode(ServiceMetaInfo serviceMetaInfo){
        for (int i = 0; i < VIRTUAL_NODES; i++) {
            int hash = getHash(serviceMetaInfo.getServiceAddress() + "#" + i);
            circle.remove(hash);
        }
    }


    /*
     * 获取目标服务器
    */

    public ServiceMetaInfo getTargetServer(Map<String, Object> requestParams) {

        int hash = getHash(requestParams.toString());

//        SortedMap<Integer, ServiceMetaInfo> tailMap = circle.tailMap(hash);
//        if(tailMap.isEmpty()){
//            return circle.firstEntry().getValue()
//        }
//
//        Integer index = tailMap.firstKey();

        Map.Entry<Integer, ServiceMetaInfo> entry = circle.ceilingEntry(hash);
        if(entry == null){
            entry = circle.firstEntry();
        }

        return entry.getValue();
    }




    /*
     * 创建虚拟节点
    */
    public void createVirtualNode(List<ServiceMetaInfo> serviceMetaInfoList){
        // 构建虚拟节点
        for (ServiceMetaInfo serviceMetaInfo: serviceMetaInfoList){
            for (int i = 0; i < VIRTUAL_NODES; i++) {
                addVirtualNode(serviceMetaInfo);
            }
        }
    }

    @Override
    public ServiceMetaInfo select(Map<String, Object> requestParams, List<ServiceMetaInfo> serviceMetaInfoList) {

        if(serviceMetaInfoList.isEmpty()) return null;

        createVirtualNode(serviceMetaInfoList);

        return getTargetServer(requestParams);
    }

    private int getHash(String key) {
        try {
            MessageDigest md5 = MessageDigest.getInstance("MD5");
            byte[] bytes = md5.digest(key.getBytes());
            int hash = ((bytes[3] & 0xFF) << 24)
                    | ((bytes[2] & 0xFF) << 16)
                    | ((bytes[1] & 0xFF) << 8)
                    | (bytes[0] & 0xFF);
            return hash & 0x7FFFFFFF;
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }
}
